import matplotlib.pyplot as plt 
import torch 
import torch.nn as nn 
import numpy as np 
import h5py
from skimage import measure, morphology
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from skimage.draw import ellipsoid
from PIL import Image
import os 
import glob
from itertools import chain
from PIL import Image

img_size = (256, 256)

def load_data(dir):
    all_path = glob.glob(dir + "*.h5")
    # print(all_path)
    data_h5 = h5py.File("./self_data_train_tumor_2d.h5", "w")
    raw_img = None 
    label_img = None 
    for each_path in all_path:
        if "t1" in each_path:
            file_h5 = h5py.File(each_path, "r")
            raw_data = file_h5["raw"][()].astype(np.float32)
            label_data = file_h5["label"][()].astype(np.float32)
            file_h5.close()
            print(raw_data.shape)
            print(label_data.shape)
            for raw, label in zip(raw_data, label_data):
                raw = Image.fromarray(raw)
                raw = raw.resize(img_size)
                label = Image.fromarray(label)
                label = label.resize(img_size)
                label = np.array(label)
                raw = np.array(raw)
                if raw_img is None and label_img is None :
                    raw_img = raw 
                    label_img = label 
                    continue
                raw_img = np.concatenate((raw_img, raw), axis=0)
                label_img = np.concatenate((label_img, label), axis=0)

    raw_img = raw_img.reshape((-1,) + img_size)
    label_img = label_img.reshape((-1,) + img_size)

    data_h5.create_dataset("raw", data=raw_img, compression="gzip")
    data_h5.create_dataset("label", data=label_img, compression="gzip")
    data_h5.close()
    return raw_img, label_img



## 构造一下2d 的输入数据 先构造 脑肿瘤的 只预测这一个
if __name__ == "__main__":
    raw, label = load_data("./self_data_train/")
    print(raw.shape)
    print(label.shape)

    # for raw1, label1 in zip(raw, label):
    #     plt.imshow(raw1, cmap="gray")
    #     plt.show()
    #     plt.imshow(label1, cmap="gray")
    #     plt.show()

    
